package com.now.presum;

import java.util.HashMap;

/**
 * 560. 和为 K 的子数组 给你一个整数数组 nums 和一个整数 k ，请你统计并返回该数组中和为 k 的连续子数组的个数。 示例 1：
 * <p>
 * 输入：nums = [1,1,1], k = 2 输出：2
 * <p>
 * 示例 2：
 * <p>
 * 输入：nums = [1,2,3], k = 3 输出：2
 */
public class Solution {

  public static int subarraySum(int[] nums, int k) {
    int n = nums.length;
    // 前缀和
    int[] preSum = new int[n + 1];
    preSum[0] = 0;
    for (int i = 1; i <= n; i++) {
      preSum[i] = preSum[i - 1] + nums[i - 1];
    }
    int res = 0;
    for (int i = 1; i <= n; i++) {
      for (int j = 0; j < i; j++) {
        if (preSum[j] == preSum[i] - k) {
          res++;
        }
      }
    }

    return res;
  }

  /**
   * 优化思路：⽤哈希表，记录前缀和的同时记录该前缀和出现的次数
    */
  public static int subarraySum2(int[] nums, int k) {
    int n = nums.length;
    HashMap<Integer, Integer> preSum = new HashMap<>();
    preSum.put(0, 1);
    int res = 0, sum0_i = 0;
    for (int i = 0; i < n; i++) {
      sum0_i += nums[i];
      // 这是我们想找的前缀和 nums[0..j]
      int sum0_j = sum0_i - k;
      // 如果前⾯有这个前缀和，则直接更新答案
      if (preSum.containsKey(sum0_j)) {
        res += preSum.get(sum0_j);
      }
      // 把前缀和 nums[0..i] 加⼊并记录出现次数
      preSum.put(sum0_i,
          preSum.getOrDefault(sum0_i, 0) + 1);
    }
    return res;
  }

  public static void main(String[] args) {
    int[] nums1 = {1, 1, 1};
    System.out.println(subarraySum(nums1, 2));
    System.out.println(subarraySum2(nums1, 2));
    int[] nums2 = {1, 2, 3};
    System.out.println(subarraySum(nums2, 3));
    System.out.println(subarraySum2(nums2, 3));
  }
}
